##########################################################################################

library(data.table)
library(optparse)
library(parallel)
library(dplyr)
library(RColorBrewer)
library(ggplot2)
library(ggsci)
library(ggpubr)

##########################################################################################

option_list <- list(
    make_option(c("--sample_list_file"), type = "character"),
    make_option(c("--signature_file"), type = "character"),
    make_option(c("--rsem_file_forstem"), type = "character"),
    make_option(c("--rsem_file_forcor"), type = "character"),
    make_option(c("--gtf_file"), type = "character"),
    make_option(c("--out_path"), type = "character")
)

if(1!=1){
    
    sample_list_file <- "~/20220915_gastric_multiple/rna_combine/analysis/config/tumor_normal.list"
    signature_file <- "~/tools/StandTools/Cibersort/SC-pcbc-stemsig.tsv"
    rsem_file_forstem <- "~/20220915_gastric_multiple/rna_combine/analysis/RSEM/CombineTPM.tsv"
    rsem_file_forcor <- "~/20220915_gastric_multiple/rna_combine/analysis/RSEM/CombineTpm.FilterLowExpression-MergeMutiSample.tsv"
    gtf_file <- "~/ref/GTF/gencode.v19.ensg_genename.txt"
    out_path <- "~/20220915_gastric_multiple/rna_combine/analysis/images/stemness"

}

parseobj <- OptionParser(option_list=option_list, usage = "usage: Rscript %prog [options]")
opt <- parse_args(parseobj)
print(opt)

sample_list_file <- opt$sample_list_file
rsem_file_forstem <- opt$rsem_file_forstem
rsem_file_forcor <- opt$rsem_file_forcor
gtf_file <- opt$gtf_file
signature_file <- opt$signature_file
out_path <- opt$out_path

dir.create(out_path , recursive = T)

##########################################################################################

info <- data.frame(fread(sample_list_file))

dat_tpm <- data.frame(fread(rsem_file_forstem))
colnames(dat_tpm) <- gsub( "X" , "" , colnames(dat_tpm))
dat_gtf <- data.frame(fread(gtf_file , header = F))
colnames(dat_gtf) <- c("gene_id" , "Hugo_Symbol")

signature <- readr::read_tsv( signature_file , col_names = F) 
signature.weight.vector <- signature$X2
names(signature.weight.vector) <- signature$X1

dat_tpm_useExpression <- data.frame(fread(rsem_file_forcor))

##########################################################################################

class_type <- c( "Normal" , "IM" , "IGC" , "DGC")
igc_class <- c("IM + IGC + DGC" , "IM + IGC")
dgc_class <- c("IM + IGC + DGC" , "IM + DGC")

im_igc_sample <- unique(info[info$Type %in% igc_class , "ID"])
im_dgc_sample <- unique(info[info$Type %in% dgc_class , "ID"])

##########################################################################################
## 按人的不同病理类型合并样本
## 若一个人同一病理类型多个样本，均中位数
dat_tpm_all <- Reduce(function(x,y)merge( x , y , by = "gene_id"),mclapply(unique(info$ID) , function(id){

    tmp_info <- subset( info , ID == id )

    result <- data.frame()
    ## 若一个人同一病理类型多个样本，均中位数
    for(class in unique(tmp_info$Class)){
        tmp_sample <- tmp_info[tmp_info$Class==class,"Tumor"]
        tmp_tpm <- dat_tpm[,c("gene_id" , tmp_sample)]

        if(ncol(tmp_tpm) > 2){
            value <- apply( tmp_tpm[,-1] , 1 , median )
        }else{
            value <- tmp_tpm[,-1]
        }
        
        result_tmp <- data.frame( gene_id = tmp_tpm$gene_id , sample = value )
        colnames(result_tmp)[2] <- paste0( id , "_" , class)

        if( nrow(result) > 0){
            result <- merge(result , result_tmp)
        }else{
            result <- result_tmp
        }
    }

    ## Normal样本
    tmp_sample <- tmp_info[tmp_info$Class==class,"Normal"]
    if(tmp_sample!="#N/A"){
        tmp_tpm <- dat_tpm[,c("gene_id" , tmp_sample)]
        value <- tmp_tpm[,-1]
        result_tmp <- data.frame( gene_id = tmp_tpm$gene_id , sample = value )
        colnames(result_tmp)[2] <- paste0( id , "_" , "Normal")

        ## 输出结果
        result <- merge(result , result_tmp)
    }
    
    result

},mc.cores=10))

## 只关注Normal、IM、IGC、DGC的表达情况
col_names <- grep( paste( c( "gene_id" , class_type) , collapse="|") , colnames(dat_tpm_all) , value = T )
dat_tpm_all <- dat_tpm_all[,col_names]

##########################################################################################

dat_tpm_all$gene_id <- sapply( strsplit(dat_tpm_all$gene_id , "[.]") , "[" , 1 )
dat_tpm_all_out <- merge( dat_gtf , dat_tpm_all)
dat_tpm_all_out <- dat_tpm_all_out[,-1]
colnames(dat_tpm_all_out)[1] <- "gene_id"
gene.expression.matrix <- dat_tpm_all_out
gene.expression.matrix <- gene.expression.matrix[ gene.expression.matrix$gene_id %in% names(signature.weight.vector),]
## 存在统一Hugo_Symbol对应多个ENSG
## 去重
gene.expression.matrix <- gene.expression.matrix[!duplicated(gene.expression.matrix$gene_id),]
rownames(gene.expression.matrix) <- gene.expression.matrix$gene_id
gene.expression.matrix <- gene.expression.matrix[,-1]

##########################################################################################
## https://github.com/BioinformaticsFMRP/PanCanStem_Web/issues/5
## Yes, they used RPKM aligned to hg19, but I would not expect an impact in
## the results using TPM or FPKM aligned to hg38.

## 计算干性评分
calculate_score <- function(signature.weight.vector, gene.expression.matrix){
  
  # Keep only common genes 
  common.genes <- intersect(names(signature.weight.vector), rownames(gene.expression.matrix))
  gene.expression.matrix <- gene.expression.matrix[common.genes, ,drop = FALSE]
  signature.weight.vector <- signature.weight.vector[common.genes]
  
  score <- apply(gene.expression.matrix, 2, function(sample) {
    cor(sample, signature.weight.vector, method = "sp", use = "complete.obs")
  })
  
  print(paste0("Min score: ",min(score)))
  print(paste0("Max score: ",max(score)))
  
  # Scale the scores to be between 0 and 1
  print(paste0("Normalized scores to be between 0 and 1"))
  score <- score - min(score)
  score.normalized <- score/max(score)
  print(paste0("Min normalized score: ",min(score.normalized)))
  print(paste0("Max normalized score: ",max(score.normalized)))
  
  return(score.normalized)
}

##########################################################################################

result <- calculate_score(signature.weight.vector,gene.expression.matrix)

sample <- sapply( strsplit(names(result)[1:length(result)] , "_") , "[" , 1)
class <- sapply( strsplit(names(result)[1:length(result)] , "_") , "[" , 2)

ratio <- as.numeric(result)
tmp_dat <- data.frame( Sample = sample , Class = class , Ratio = ratio )

## Normal->IM->IGC看基因的表达变化情况
tmp_dat_igc <- subset( tmp_dat , Sample %in% im_igc_sample )
tmp_dat_igc <- subset( tmp_dat_igc , Class %in% c( "Normal" , "IM" , "IGC" ))
tmp_dat_igc$Class <- factor( tmp_dat_igc$Class , levels = c( "Normal" , "IM" , "IGC" ) , order = T )
tmp_dat_igc$Type <- "IM + IGC"

## Normal->IM->DGC看基因的表达变化情况
tmp_dat_dgc <- subset( tmp_dat , Sample %in% im_dgc_sample )
tmp_dat_dgc <- subset( tmp_dat_dgc , Class %in% c( "Normal" , "IM" , "DGC" ))
tmp_dat_dgc$Class <- factor( tmp_dat_dgc$Class , levels = c( "Normal" , "IM" , "DGC" ) , order = T )
tmp_dat_dgc$Type <- "IM + DGC"

## Normal->IM->GC看基因的表达变化情况
tmp_dat_gc <- tmp_dat
tmp_dat_gc$Class <- ifelse( tmp_dat_gc$Class %in% c("IGC" , "DGC") , "GC" , tmp_dat_gc$Class )
tmp_dat_gc$Class <- factor( tmp_dat_gc$Class , levels = c( "Normal" , "IM" , "GC" ) , order = T )
tmp_dat_gc$Type <- "IM + GC"
## 一个既有IGC又有DGC，合并
tmp_dat_gc <- tmp_dat_gc %>%
group_by(Sample , Class) %>%
summarize( Ratio = median(Ratio) , Type = unique(Type) )

dat_plot <- rbind( tmp_dat_igc , tmp_dat_dgc , tmp_dat_gc )
dat_plot$Type <- factor( dat_plot$Type , levels = c("IM + GC" , "IM + IGC" , "IM + DGC") )


##########################################################################################
out_name <- paste0( out_path , "/StemScore.MutipleStage.pdf" )

my_comparisons_1 <- list( 
    c(1, 2), c(1, 3) , 
    c(2, 3) )

plot <- ggplot( dat_plot , aes( x = Class , y = Ratio , color = Class ) ) +
    geom_line( aes( group = Sample ) , size = 0.4 , color = "gray" ) +  ## 配对样本加线
    geom_boxplot(alpha =1 , outlier.size=0 , size = 0.9 , width = 0.6) +
    geom_jitter(position = position_jitter(0.17) , size = 1 , alpha = 0.7) +
    facet_grid(.~Type,space='free_x',scales='free_x') +
    scale_fill_npg()+
    scale_color_npg()+
    xlab(NULL) +
    ylab("Stem score")+
    theme_bw() +
    stat_compare_means(comparisons = my_comparisons_1) +
    theme(panel.background = element_blank(),#设置背影为白色#清除网格线
        legend.position ='none',
        legend.title = element_blank() ,
        panel.grid.major=element_line(colour=NA),
        plot.title = element_text(size = 8,color="black",face='bold'),
        legend.text = element_text(size = 10,color="black",face='bold'),
        axis.text.y = element_text(size = 10,color="black",face='bold'),
        axis.title.x = element_text(size = 10,color="black",face='bold'),
        axis.title.y = element_text(size = 10,color="black",face='bold'),
        axis.ticks.x = element_blank(),
        axis.text.x = element_text(size = 10,color="black",face='bold') ,
        axis.line = element_line(size = 0.5)) 
ggsave(file=out_name,plot=plot,width=6,height=4)

out_name <- paste0( out_path , "/StemScore.MutipleStage.tsv" )
write.table( dat_plot , out_name , row.names = F , sep = "\t" , quote = F )

##########################################################################################
out_name <- paste0( out_path , "/StemScore.MutipleStage.oneImage.pdf" )

my_comparisons_1 <- list( 
    c(1, 2) , c(1, 3) , c(1, 4) , 
    c(2, 3) , c(2, 4) ,
    c(3, 4)  
    )

tmp_dat$Class <- factor( tmp_dat$Class , levels = c( "Normal" , "IM" , "IGC" , "DGC" ) , order = T )
plot <- ggplot( tmp_dat , aes( x = Class , y = Ratio , color = Class ) ) +
    geom_line( aes( group = Sample ) , size = 0.4 , color = "gray" ) +  ## 配对样本加线
    geom_boxplot(alpha =1 , outlier.size=0 , size = 0.9 , width = 0.6) +
    geom_jitter(position = position_jitter(0.17) , size = 1 , alpha = 0.7) +
    scale_fill_npg()+
    scale_color_npg()+
    xlab(NULL) +
    ylab("Stem score")+
    theme_bw() +
    stat_compare_means(comparisons = my_comparisons_1) +
    theme(panel.background = element_blank(),#设置背影为白色#清除网格线
        legend.position ='none',
        legend.title = element_blank() ,
        panel.grid.major=element_line(colour=NA),
        plot.title = element_text(size = 8,color="black",face='bold'),
        legend.text = element_text(size = 10,color="black",face='bold'),
        axis.text.y = element_text(size = 10,color="black",face='bold'),
        axis.title.x = element_text(size = 10,color="black",face='bold'),
        axis.title.y = element_text(size = 10,color="black",face='bold'),
        axis.ticks.x = element_blank(),
        axis.text.x = element_text(size = 10,color="black",face='bold') ,
        axis.line = element_line(size = 0.5)) 
ggsave(file=out_name,plot=plot,width=4,height=5)

out_name <- paste0( out_path , "/StemScore.MutipleStage.oneImage.tsv" )
write.table( dat_plot , out_name , row.names = F , sep = "\t" , quote = F )

##########################################################################################
## 评估基因表达和干细胞评分的关系
dat_stemscore <- data.frame(stem_score = result , ID = names(result))

dat_cor <- c()
for( gene in dat_tpm_useExpression$Hugo_Symbol ){
    print(gene)
    tmp_tpm <- subset( dat_tpm_useExpression , Hugo_Symbol == gene )
    tmp_tpm <- t(tmp_tpm[,-c(1:2)])
    colnames(tmp_tpm) <- "TPM"
    tmp_tpm <- data.frame(tmp_tpm)
    tmp_tpm$ID <- rownames(tmp_tpm)

    tmp_tpm <- merge(tmp_tpm , dat_stemscore , by = "ID" )
    tmp_tpm$Hugo_Symbol <- gene

    cor <- as.numeric(cor.test( tmp_tpm$TPM , tmp_tpm$stem_score , method = "spearman" )$estimate)
    p <- as.numeric(cor.test( tmp_tpm$TPM , tmp_tpm$stem_score , method = "spearman"  )$p.value)

    tmp_dat <- data.frame( Hugo_Symbol = gene , cor = cor , p = p  )

    dat_cor <- rbind( dat_cor , tmp_dat )
}

out_name <- paste0( out_path , "/StemScore_Expression.tsv" )
dat_cor$q <- p.adjust( dat_cor$p , method  = "fdr" )
write.table( dat_cor , out_name , row.names = F , quote = F , sep = "\t" )